#include <iostream>
#include <string>
#include <map>
#include <winsock2.h>
#include <ws2tcpip.h>
#pragma comment(lib, "ws2_32.lib")

#define DHCP_SERVER_PORT 67
#define DHCP_CLIENT_PORT 68
#define FIXED_IP "192.168.1.2"
#define SUBNET_MASK "255.255.255.0"
#define GATEWAY "192.168.1.1"
#define DNS_SERVER "8.8.8.8"

// DHCP报文结构
#pragma pack(push, 1)
struct dhcp_packet
{
    uint8_t op;
    uint8_t htype;
    uint8_t hlen;
    uint8_t hops;
    uint32_t xid;
    uint16_t secs;
    uint16_t flags;
    uint32_t ciaddr;
    uint32_t yiaddr;
    uint32_t siaddr;
    uint32_t giaddr;
    uint8_t chaddr[16];
    char sname[64];
    char file[128];
    uint32_t magic;
    uint8_t options[308];
};
#pragma pack(pop)

// DHCP选项类型
enum dhcp_option
{
    OPT_PAD = 0,
    OPT_SUBNET_MASK = 1,
    OPT_ROUTER = 3,
    OPT_DNS_SERVER = 6,
    OPT_REQUESTED_IP = 50,
    OPT_LEASE_TIME = 51,
    OPT_MESSAGE_TYPE = 53,
    OPT_SERVER_ID = 54,
    OPT_END = 255
};

// DHCP消息类型
enum dhcp_message_type
{
    DHCPDISCOVER = 1,
    DHCPOFFER = 2,
    DHCPREQUEST = 3,
    DHCPACK = 5
};

// 添加DHCP选项
void add_option(uint8_t *options, int &offset, uint8_t code, const void *data, uint8_t len)
{
    options[offset++] = code;
    options[offset++] = len;
    memcpy(&options[offset], data, len);
    offset += len;
}

// 处理DHCP DISCOVER
void handle_discover(SOCKET sockfd, sockaddr_in client_addr, dhcp_packet *packet)
{
    char mac[18];
    snprintf(mac, sizeof(mac), "%02X:%02X:%02X:%02X:%02X:%02X",
             packet->chaddr[0], packet->chaddr[1], packet->chaddr[2],
             packet->chaddr[3], packet->chaddr[4], packet->chaddr[5]);

    std::string allocated_ip = FIXED_IP;
    std::cout << "Assigning fixed IP: " << allocated_ip << " to " << mac << std::endl;

    dhcp_packet offer = {0};
    offer.op = 0x02; // BOOTREPLY
    offer.htype = 1; // Ethernet
    offer.hlen = 6;
    offer.xid = packet->xid;
    offer.yiaddr = inet_addr(allocated_ip.c_str());
    offer.siaddr = inet_addr(GATEWAY);
    offer.magic = htonl(0x63825363); // Magic cookie

    int opt_offset = 0;
    uint8_t msg_type = DHCPOFFER;
    add_option(offer.options, opt_offset, OPT_MESSAGE_TYPE, &msg_type, 1);

    uint32_t lease_time = htonl(86400); // 24小时
    add_option(offer.options, opt_offset, OPT_LEASE_TIME, &lease_time, 4);

    uint32_t server_id = inet_addr(GATEWAY);
    add_option(offer.options, opt_offset, OPT_SERVER_ID, &server_id, 4);

    uint32_t subnet_mask = inet_addr(SUBNET_MASK);
    add_option(offer.options, opt_offset, OPT_SUBNET_MASK, &subnet_mask, 4);

    uint32_t router = inet_addr(GATEWAY);
    add_option(offer.options, opt_offset, OPT_ROUTER, &router, 4);

    uint32_t dns = inet_addr(DNS_SERVER);
    add_option(offer.options, opt_offset, OPT_DNS_SERVER, &dns, 4);

    offer.options[opt_offset++] = OPT_END;

    sendto(sockfd, (char *)&offer, sizeof(offer), 0,
           (sockaddr *)&client_addr, sizeof(client_addr));
}

int main()
{
    WSADATA wsa;
    if (WSAStartup(MAKEWORD(2, 2), &wsa) != 0)
    {
        std::cerr << "WSAStartup failed" << std::endl;
        return 1;
    }

    SOCKET sockfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sockfd == INVALID_SOCKET)
    {
        std::cerr << "Socket creation error" << std::endl;
        return 1;
    }

    int optval = 1;
    setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, (char *)&optval, sizeof(optval));

    sockaddr_in server_addr{};
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(DHCP_SERVER_PORT);
    server_addr.sin_addr.s_addr = INADDR_ANY;

    if (bind(sockfd, (sockaddr *)&server_addr, sizeof(server_addr)) == SOCKET_ERROR)
    {
        std::cerr << "Bind failed" << std::endl;
        closesocket(sockfd);
        return 1;
    }

    std::cout << "DHCP Server started..." << std::endl;

    while (true)
    {
        dhcp_packet packet;
        sockaddr_in client_addr{};
        int addr_len = sizeof(client_addr);

        int len = recvfrom(sockfd, (char *)&packet, sizeof(packet), 0,
                           (sockaddr *)&client_addr, &addr_len);
        if (len == SOCKET_ERROR)
        {
            std::cerr << "Receive error" << std::endl;
            continue;
        }

        uint8_t *options = packet.options;
        uint8_t message_type = 0;
        while (*options != OPT_END)
        {
            if (*options == OPT_MESSAGE_TYPE)
            {
                message_type = *(options + 2);
                break;
            }
            options += options[1] + 2;
        }

        switch (message_type)
        {
        case DHCPDISCOVER:
            handle_discover(sockfd, client_addr, &packet);
            break;
        case DHCPREQUEST:
            // 处理REQUEST
            break;
        }
    }

    closesocket(sockfd);
    WSACleanup();
    return 0;
}